from utils.pdf_processor import PDFProcessor
from utils.image_processor import ImageProcessor
from utils.text_processor import TextProcessor
from models.rag_model import RAGModel
import os
from config import *

def main():
    # 初始化处理器
    pdf_processor = PDFProcessor(PDF_UPLOAD_DIR)
    image_processor = ImageProcessor()
    text_processor = TextProcessor()
    rag_model = RAGModel(COLLECTION_NAME)
    
    # 处理PDF文件
    pdf_path = r"D:\test_document.pdf"
    processed_data = pdf_processor.process_pdf(pdf_path)
    
    # 生成向量
    embeddings = []
    for item in processed_data:
        if item["type"] == "text":
            emb = text_processor.get_text_embedding(item["content"])
        else:  # image
            emb = image_processor.get_image_embedding(item["image"])
        embeddings.append(emb)
    
    # 存储到Milvus
    rag_model.add_documents(processed_data, embeddings)
    
    # 问答示例
    def answer_question(question: str):
        # 获取问题向量
        question_embedding = text_processor.get_text_embedding(question)
        
        # 搜索相关文档
        relevant_docs = rag_model.search(question_embedding)
        
        # 构建提示词
        context = "\n".join([f"Page {doc['page']}: {doc['content']}" for doc in relevant_docs])
        prompt = f"""基于以下文档内容回答问题：
        
文档内容：
{context}

问题：{question}

请结合文档内容回答问题。如果文档中没有相关信息，请说明无法回答。"""

        # 调用千问模型
        # 这里需要实现调用千问API的代码
        # response = call_qwen_api(prompt)
        
        return prompt

    print(answer_question("2022年营收是多少？"))


if __name__ == "__main__":
    main() 